Can you explain how you are handling broadcasting with vmap? The code looks very confusing.

The vmap logic here is just normalizing all inputs into a canonical batched layout, then reusing the existing output_shape() logic unchanged.

1. If all inputs are unbatched (axes == -1), it preserves the old path exactly.
2. Otherwise it computes ndim = max(input.ndim() + (axis == -1)), so unbatched inputs reserve one extra dimension for the mapped axis.
3. Each input is left-padded with leading size-1 dimensions to reach ndim.
4. It then aligns the mapped axis of every input to a single physical position to_ax:
   • if an input is already vmapped, its existing mapped axis is moved to to_ax
   • if an input is unvmapped, the inserted singleton dimension serves as its mapped axis
5. ignore_axes_ is defined relative to the original unbatched layout, so after inserting/aligning the mapped axis it has to be remapped into the normalized layout. The code converts each ignored axis to a positive index in the unbatched view, shifts it if the mapped axis was inserted before it, applies the left-padding offset, then converts it back to the negative indexing convention expected by output_shape().

After that, output_shape(aligned_inputs, remapped_ignore_axes) computes the same broadcast result as before, just on the normalized batched representation, and the output mapped axis is to_ax.

I can also add a few comments in the implementation to make that flow easier to follow.

Well that is just a rephrasing of what the code does. To be honest I don't know if you actually understood the algorithm and explained it to me, or just pasted what AI said. I'm leaving this to @angeloskath to judge.